
#include "Constraint.h"

#include "Creature.h"

#include "BasicCreature.h"
#include "QuadrupedCreature.h"
#include "HumanCreature.h"

#include "BulletSimulator.h"
#include "NeuralNetwork.h"

#include "MathUtil.h"
#include "BasicDebug.h"
#include "HalfHumanCreature.h"

#include <QtCore/QString>

namespace jcrada {

    Creature::Creature(int body_parts, int constraints)
    :
    _name(""),
    _number_of_body_parts(body_parts),
    _number_of_constraints(constraints),
    _body_parts(new BodyPart*[body_parts]),
    _constraints(new Constraint*[constraints]),
    _neural_network(NULL),
    _initial_position(btVector3(0.0, 0.0, 0.0)),
    _final_position(btVector3(0.0, 0.0, 0.0)),
    _fitness(-1.0) {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            setBodyPart(i, NULL);
        }
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            setConstraint(i, NULL);
        }
    }

    Creature::Creature(int body_parts, int constraints,
            int hidden_layers, int neurons_per_layer)
    :
    _name(""),
    _number_of_body_parts(body_parts),
    _number_of_constraints(constraints),
    _body_parts(new BodyPart*[body_parts]),
    _constraints(new Constraint*[constraints]),
    _neural_network(NULL),
    _initial_position(btVector3(0.0, 0.0, 0.0)),
    _final_position(btVector3(0.0, 0.0, 0.0)),
    _fitness(0.0) {

        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            setBodyPart(i, new BodyPart());
            getBodyPart(i).setId(i);
            getBodyPart(i).setName("Body Part #" + TO_STRING(i));
        }
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            setConstraint(i, new Constraint());
            getConstraint(i).setId(i);
            getConstraint(i).setName("Constraint #" + TO_STRING(i));
        }

        _neural_network = new NeuralNetwork(4 * getNumberOfBodyParts(),
                hidden_layers, neurons_per_layer, 3 * (getNumberOfBodyParts() - 1));
        getNeuralNetwork().getInputLayer().setActivationFunction(Neuron::AF_NONE);
        for (int i = 1; i < getNeuralNetwork().getNumberOfLayers(); ++i) {
            getNeuralNetwork().getLayer(i).setActivationFunction(Neuron::AF_TANH);
        }
    }

    Creature::Creature(int body_parts, int constraints, const NeuralNetwork & neural_network)
    :
    _name(""),
    _number_of_body_parts(body_parts),
    _number_of_constraints(constraints),
    _body_parts(new BodyPart*[body_parts]),
    _constraints(new Constraint*[constraints]),
    _neural_network(NULL),
    _initial_position(btVector3(0.0, 0.0, 0.0)),
    _final_position(btVector3(0.0, 0.0, 0.0)),
    _fitness(0.0) {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            setBodyPart(i, new BodyPart());
            getBodyPart(i).setName("Body Part #" + TO_STRING(i));
        }
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            setConstraint(i, new Constraint());
            getConstraint(i).setName("Constraint #" + TO_STRING(i));
        }

        setNeuralNetwork(neural_network.clone());
    }

    Creature::~Creature() {
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            delete _constraints[i];
        }
        delete [] _constraints;

        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            delete _body_parts[i];
        }
        delete [] _body_parts;

        delete _neural_network;
    }

    BodyPart & Creature::getRoot() const {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            if (getBodyPart(i).isRoot()) {
                return getBodyPart(i);
            }
        }
        assert(0);
    }

    int Creature::getIndex(const BodyPart & body) const {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            if (&getBodyPart(i) == &body) {
                return i;
            }
        }
        return -1;
    }

    int Creature::getIndex(const btRigidBody * body) const {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            if (getBodyPart(i).getRigidBody() == body) {
                return i;
            }
        }
        return -1;
    }

    int Creature::getIndex(const Constraint & constraint) const {
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            if (&getConstraint(i) == &constraint) {
                return i;
            }
        }
        return -1;
    }

    //    void Creature::set( const btTransform & at ) {
    //        btTransform offset;
    //        offset.setIdentity();
    //        double max_height = 0.0;
    //        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
    //            max_height += getBodyPart(i).getSizeY();
    //        }
    //        offset.setOrigin(btVector3(0.0, max_height, 0.0));
    //        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
    //            getBodyPart(i).set(at * offset);
    //        }
    //        for (int i = 0; i < getNumberOfConstraints(); ++i) {
    //            getConstraint(i).set(*this);
    //        }
    //        calculateMaxTorque();
    //    }

    void Creature::set(const btTransform & at) {
        btTransform offset;
        offset.setIdentity();
        double avg_distance = 0.0;
        double max_height = getMaxHeight();
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            avg_distance += (getBodyPart(i).getSizeX() + getBodyPart(i).getSizeY() +
                    getBodyPart(i).getSizeZ()) / 3.0;
        }
        avg_distance *= 2;
        avg_distance /= getNumberOfBodyParts();
        offset.setOrigin(btVector3(0.0, avg_distance, 0.0));
        float circle_fraction = 2 * M_PI / getNumberOfBodyParts();
        float next_fraction = 0.0;
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            offset.setIdentity();
            if (getBodyPart(i).isRoot()) {
                offset.setOrigin(btVector3(0, max_height, 0));
            } else {
                offset.setOrigin(btVector3(cos(next_fraction) * avg_distance, max_height, sin(next_fraction) * avg_distance));
            }
            getBodyPart(i).set(at * offset);
            next_fraction += circle_fraction;
        }
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            getConstraint(i).set(*this);
        }
        calculateMaxTorque();
    }

    void Creature::reset() {
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            getConstraint(i).reset();
        }
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            getBodyPart(i).reset();
        }
    }

    void Creature::scale(const btScalar factor) {
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            getBodyPart(i).scaleSizeAndMass(factor);
        }
    }

    void Creature::loadConnectionTransform(const BodyPart& body, const Constraint& constraint, btTransform & trans) {
        if (body.getId() == constraint.getIdBodyA()) {
            Constraint::locateConnectionOrigin(body, constraint.getConnectionA(), trans);
        } else {

            assert(body.getId() == constraint.getIdBodyB());
            Constraint::locateConnectionOrigin(body, constraint.getConnectionB(), trans);
        }
        trans = body.getRigidBody()->getCenterOfMassTransform() * trans;
    }

    Constraint * Creature::getConstraint(const BodyPart& a, const BodyPart & b) const {
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            if ((getConstraint(i).getIdBodyA() == a.getId() && getConstraint(i).getIdBodyB() == b.getId()) ||
                    (getConstraint(i).getIdBodyA() == b.getId() && getConstraint(i).getIdBodyB() == a.getId())) {
                return &getConstraint(i);
            }
        }
        return NULL;
    }

    void Creature::calculateMaxTorque() {
        double gravity = Math::abs(BulletSimulator::getSingleton().getWorld().getGravity().y());
        int terminal[getNumberOfBodyParts()];
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            terminal[i] = 0;
            getBodyPart(i).setMaxTorque(0.0);
        }
        //Determinar nodos terminales (sum == 1)
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            assert(getConstraint(i).getIdBodyA() < getNumberOfBodyParts() &&
                    getConstraint(i).getIdBodyA() >= 0);
            assert(getConstraint(i).getIdBodyB() < getNumberOfBodyParts() &&
                    getConstraint(i).getIdBodyB() >= 0);
            //            if (!getBodyPart(getConstraint(i).getIdBodyA()).isRoot()) {
            ++terminal[getConstraint(i).getIdBodyA()];
            //            }
            //            if (!getBodyPart(getConstraint(i).getIdBodyB()).isRoot()) {
            ++terminal[getConstraint(i).getIdBodyB()];
            //            }
        }
        for (int e = 0; e < getNumberOfBodyParts(); ++e) {
            if (terminal[e] != 1) continue;
            int count = 0;
            BodyPart * body = &getBodyPart(e);
            std::vector<BodyPart*> extremities;
            while (body != &getRoot() && count++ < getNumberOfBodyParts()) {
                extremities.push_back(body);
                body = &getBodyPart(body->getParentId());
            }
            extremities.push_back(body);

            double arm_length = 0.0;
            btTransform trans;
            btVector3 center_of_mass;
            btVector3 connection;
            // A  0  B  1  C  2  Root  3
            // |-----|-----|-----|-----|
            for (int i = extremities.size() - 2; i >= 0; --i) {
                //                DEBUG(extremity[i]->getName() + " ; " + extremity[i+1]->getName());
                Constraint* constraint = getConstraint(*extremities[i], *extremities[i + 1]);
                assert(constraint);
                center_of_mass = extremities[i]->getRigidBody()->getCenterOfMassPosition();
                loadConnectionTransform(*extremities[i], *constraint, trans);
                connection = btVector3(trans.getOrigin().x(), trans.getOrigin().y(), trans.getOrigin().z());
                arm_length = (center_of_mass - connection).length();
                //                BDEBUG("arm=" + TO_STRING(arm_length));
                extremities[i]->addMaxTorque(arm_length * extremities[i]->getMass() * gravity);

                for (int j = i - 1; j >= 0; --j) {
                    constraint = getConstraint(*extremities[j], *extremities[i]);
                    if (!constraint) {
                        //                        DEBUG("NOT FOUND: " + extremity[i]->getName() + " ; " + extremity[i + 1]->getName());

                        continue;
                    }
                    //                    DEBUG("FOUND (" + constraint->getName() + "): " + extremity[i]->getName() + " ; " + extremity[i + 1]->getName());
                    assert(constraint);
                    center_of_mass = extremities[i]->getRigidBody()->getCenterOfMassPosition();
                    loadConnectionTransform(*extremities[i], *constraint, trans);
                    connection = btVector3(trans.getOrigin().x(), trans.getOrigin().y(), trans.getOrigin().z());
                    arm_length += (center_of_mass - connection).length();

                    center_of_mass = extremities[j]->getRigidBody()->getCenterOfMassPosition();
                    loadConnectionTransform(*extremities[j], *constraint, trans);
                    connection = btVector3(trans.getOrigin().x(), trans.getOrigin().y(), trans.getOrigin().z());
                    arm_length += (center_of_mass - connection).length();

                    extremities[i]->addMaxTorque(arm_length * extremities[j]->getMass() * gravity);
                }
            }
        }
    }

    void Creature::onHitTheGround() {
        setInitialPosition(getRoot().getRigidBody()->getCenterOfMassPosition());
    }

    void Creature::beforeRemove() {
        setFinalPosition(getRoot().getRigidBody()->getCenterOfMassPosition());
    }

    void Creature::onStep() {
        btTransform trans;
        int index = 0;
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            getBodyPart(i).getRigidBody()->getMotionState()->getWorldTransform(trans);
            getNeuralNetwork().getInputLayer().getNeuron(index++).setInput(0, trans.getRotation().w());
            getNeuralNetwork().getInputLayer().getNeuron(index++).setInput(0, trans.getRotation().x());
            getNeuralNetwork().getInputLayer().getNeuron(index++).setInput(0, trans.getRotation().y());
            getNeuralNetwork().getInputLayer().getNeuron(index++).setInput(0, trans.getRotation().z());
        }
        getNeuralNetwork().process();
        btVector3 torque;
        index = 0;
        double mx = 0.0, my = 0.0, mz = 0.0, sum_m = 0.0;
        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            if (getBodyPart(i).isRoot()) {
                continue;
            }
            getBodyPart(i).getRigidBody()->activate(true);
            mx = getNeuralNetwork().getOutputLayer().getNeuron(index++).output();
            my = getNeuralNetwork().getOutputLayer().getNeuron(index++).output();
            mz = getNeuralNetwork().getOutputLayer().getNeuron(index++).output();
            //            sum_m = Math::abs(mx) + Math::abs(mz);
            //            torque.setValue(mx * getBodyPart(i).getMaxTorque() / sum_m,
            //                    my * 0.1 * getBodyPart(i).getMaxTorque(),
            //                    mz * getBodyPart(i).getMaxTorque() / sum_m);
            torque.setValue(mx * getBodyPart(i).getMaxTorque(),
                    my * 0.1 * getBodyPart(i).getMaxTorque(),
                    mz * getBodyPart(i).getMaxTorque());
            getBodyPart(i).getRigidBody()->applyTorque(torque);
        }
    }

    Creature * Creature::clone() const {
        Creature* result = new Creature(getNumberOfBodyParts(), getNumberOfConstraints(),
                getNeuralNetwork());
        for (int i = 0; i < result->getNumberOfBodyParts(); ++i) {
            result->setBodyPart(i, getBodyPart(i).clone());
        }
        for (int i = 0; i < result->getNumberOfConstraints(); ++i) {
            result->setConstraint(i, getConstraint(i).clone());
        }
        result->setFinalPosition(getFinalPosition());
        result->setInitialPosition(getInitialPosition());
        result->setFitness(getFitness());
        result->setName(getName());

        return result;
    }
    //-----------------------------------------------------------
    //XML
    //-----------------------------------------------------------

    Creature * Creature::fromXml(const QDomElement& xml, bool* ok) {
        int body_parts = xml.attribute("body-parts", "0").toInt(ok);
        int constraints = xml.attribute("constraints", "0").toInt(ok);
        Creature* result = new Creature(body_parts, constraints);
        result->setName(xml.attribute("name", "").toStdString());
        result->setFitness(xml.attribute("fitness", "0").toDouble(ok));
        QDomNode child = xml.firstChild();
        int ix_body = 0, ix_constraint = 0;
        while (!child.isNull()) {
            QDomElement e = child.toElement();
            if (e.tagName() == "BodyPart") {
                assert(ix_body < result->getNumberOfBodyParts());
                result->setBodyPart(ix_body, BodyPart::fromXml(e, ok));
                ++ix_body;
            } else if (e.tagName() == "Constraint") {
                assert(ix_constraint < result->getNumberOfConstraints());
                result->setConstraint(ix_constraint, Constraint::fromXml(e, ok));
                ++ix_constraint;
            } else if (e.tagName() == "NeuralNetwork") {
                result->setNeuralNetwork(NeuralNetwork::fromXml(e, ok));
            }
            child = child.nextSibling();
        }
        result->set();

        return result;
    }

    Creature * Creature::fromXml(const std::string& xml, bool* ok) {
        QDomDocument doc;
        QString * error_msg = new QString();
        int* error_line = new int();
        int* error_col = new int();
        *ok = doc.setContent(QString(xml.c_str()), error_msg, error_line, error_col);
        if (!*ok) {
            BDEBUG("Error: " + error_msg->toStdString() +
                    " at [" + TO_STRING(*error_line) + " , " +
                    TO_STRING(*error_col) + "]");
        }
        delete error_msg;
        delete error_line;
        delete error_col;
        if (!*ok) {
            return NULL;
        }
        QDomElement root = doc.firstChild().toElement();

        return fromXml(root, ok);
    }

    void Creature::toXml(QDomDocument& xml, QDomElement * parent) const {
        QDomElement creature = xml.createElement("Creature");

        //        creature.setAttribute("id", getId());
        creature.setAttribute("name", QString(getName().c_str()));
        //        creature.setAttribute("type", QString::number(getType()));
        creature.setAttribute("body-parts", getNumberOfBodyParts());
        creature.setAttribute("constraints", getNumberOfConstraints());
        creature.setAttribute("hidden-layers", getNeuralNetwork().getNumberOfLayers() - 2);
        creature.setAttribute("neurons-per-layer", getNeuralNetwork().getLayer(1).getNumberOfNeurons());
        creature.setAttribute("initial-position", QString((
                TO_STRING(getInitialPosition().x()) + " ; " +
                TO_STRING(getInitialPosition().y()) + " ; " +
                TO_STRING(getInitialPosition().z())).c_str()));
        creature.setAttribute("final-position", QString((
                TO_STRING(getFinalPosition().x()) + " ; " +
                TO_STRING(getFinalPosition().y()) + " ; " +
                TO_STRING(getFinalPosition().z())).c_str()));
        creature.setAttribute("fitness", getFitness());

        for (int i = 0; i < getNumberOfBodyParts(); ++i) {
            getBodyPart(i).toXml(xml, &creature);
        }
        for (int i = 0; i < getNumberOfConstraints(); ++i) {
            getConstraint(i).toXml(xml, &creature);
        }
        getNeuralNetwork().toXml(xml, &creature);
        if (parent) {
            parent->appendChild(creature);
        } else {

            xml.appendChild(creature);
        }
    }

    std::string Creature::toXml() const {
        QDomDocument doc;
        this->toXml(doc);

        return doc.toString().toStdString();
    }

    int Creature::main(int argc, char** argv) {
        (void) argc;
        (void) argv;
        BasicCreature creature(1);

        QDomDocument doc;
        creature.toXml(doc);
        QString xml = doc.toString();
        //        std::cout << xml.toStdString() << std::endl;
        BDEBUG(xml.toStdString());

        doc.clear();
        QString * error_msg = new QString();
        int* error_line = new int();
        int* error_col = new int();
        doc.setContent(xml, error_msg, error_line, error_col);
        std::cout << "Error: " << error_msg->toStdString() << " at [" << *error_line << " , "
                << *error_col << "]" << std::endl;
        delete error_msg;
        delete error_line;
        delete error_col;

        QDomElement root = doc.firstChild().toElement();
        //        creature.fromXml(root);
        doc.clear();
        creature.toXml(doc);
        xml = doc.toString();
        BDEBUG("******************************");
        BDEBUG(xml.toStdString());
        std::cout << xml.toStdString() << std::endl;
        return 0;
    }
}
